import time

from daily_tasks import DailyTasks
from fight import Fight
from repair_in import RepairIn
from switch import Switch

import torch
from ultralytics import YOLO


class Controller:

    def __init__(self, position_grab, map_name):

        self.region = (0, 0, 960, 570)

        self.map_name = map_name

        self.position_grab = position_grab

        self.model = None

        self.is_stop = False

    def start(self):
        # =====================加载AI模型==============================
        self.model = YOLO(f'model/{self.map_name}.pt', task='detect')
        # 检查是否有可用的GPU
        use_cuda = torch.cuda.is_available()
        device = torch.device("cuda:0" if use_cuda else "cpu")
        print(f"当前设备：{device}")
        # 将模型加载到设备上
        self.model.to(device)
        print(f"[Controller]:加载AI模型")
        time.sleep(2)

        # =====================日常任务=====================
        DailyTasks.daily_main(self.region)

        # 当前模式：0、检查修理装备； 1、关卡战斗; 2、校验体力，选择角色
        current_mode = 2

        # 战斗场景
        fight = Fight(self.map_name, self.position_grab, self.model)

        # 切换角色
        switch = Switch(self.position_grab, self.map_name)

        # 修理装备
        repair = RepairIn(self.position_grab)

        while True:
            # 执行流程：2[校验体力，选择角色] -> 1[关卡战斗] -> 0[检查修理装备]...-> 1[关卡战斗] -> 0[检查修理装备] ->2[校验体力，选择角色]

            if current_mode == 0:
                is_continue, current_mode = repair.repair_main()
                if is_continue:
                    continue

            elif current_mode == 1:
                is_continue, current_mode = fight.fight_main()
                if is_continue:
                    continue

            elif current_mode == 2:
                is_continue, current_mode = switch.switch_main()
                if is_continue:
                    if switch.is_Switch:
                        # 更换角色之后，切换刷图类型
                        break
                    else:
                        continue
                else:
                    self.is_stop = True
                    break
